import numpy as np
import matplotlib.pyplot as plt


def CreateTestData(origin_x, origin_y, noise_point=1):
    test_x = []
    test_y = []
    for x, y in zip(origin_x, origin_y):
        for i in range(len(x)):
            if noise_point == 2:
                for j in range(i + 1, len(x)):
                    new_x = x.copy()
                    new_x[i] = 1 if new_x[i] == 0 else 0
                    new_x[j] = 1 if new_x[j] == 0 else 0
                    test_x.append(new_x)
                    test_y.append(y)
            else:
                new_x = x.copy()
                new_x[i] = 1 if new_x[i] == 0 else 0
                test_x.append(new_x)
                test_y.append(y)
    return np.array(test_x), np.array(test_y)


class Perceptron(object):

    def __init__(self, input_dim=25, output_dim=3):
        self.weight = np.zeros([input_dim + 1, output_dim])

    def forward(self, input):
        self.input = np.hstack([1, input])
        self.out = np.sign(np.dot(self.input, self.weight))
        return self.out

    def backward(self, label, lr):
        self.weight += lr * np.matmul(self.input.reshape([-1, 1]),
                                      (label - self.out).reshape([1, -1]))


def train(model, X, Y, n_iter, lr):
    losses = []
    for i in range(n_iter):
        loss = 0
        for x, y in zip(X, Y):
            pred = model.forward(x)
            loss += np.sum((y - pred)**2)
            model.backward(y, lr)
        loss /= X.shape[0]
        losses.append(loss)
        print("iter:%d loss:%.10f" % (i, loss))
    return losses


def test(model, test_x, test_y):
    loss = 0
    acc = 0
    for i in range(test_x.shape[0]):
        pred = model.forward(test_x[i])
        loss += np.sum((test_y[i] - pred)**2)
        if np.argmax(pred) == np.argmax(test_y[i]):
            acc += 1
    loss /= test_x.shape[0]
    acc /= test_x.shape[0]
    print("test_loss:%f\ntest_acc:%f" % (loss, acc))


C = [[0, 1, 1, 1, 0], [1, 0, 0, 0, 1], [1, 0, 0, 0, 0], [1, 0, 0, 0, 1],
     [0, 1, 1, 1, 0]]
H = [[1, 0, 0, 0, 1], [1, 0, 0, 0, 1], [1, 1, 1, 1, 1], [1, 0, 0, 0, 1],
     [1, 0, 0, 0, 1]]
L = [[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0],
     [1, 1, 1, 1, 1]]
x = np.array([C, H, L]).reshape([-1, 25])
y = np.array([[1, -1, -1], [-1, 1, -1], [-1, -1, 1]])

n_iter = 20
lr = 5
model = Perceptron()
noise1_x, noise1_y = CreateTestData(x, y)
noise2_x, noise2_y = CreateTestData(x, y, 2)

# train with origin data
print("train with origin data")
plt.xlabel('iter')
plt.ylabel('loss')


losses = train(model, x, y, n_iter, lr)
plt.plot(range(n_iter), losses, "-", 'red')
plt.savefig("train1.jpg")
plt.cla()
print("test in data with a noise point")
test(model, noise1_x, noise1_y)
print("test in data with two noise point")
test(model, noise2_x, noise2_y)

# train with origin data and data with a noise point
model = Perceptron()
print("train with origin data and data with a noise point")
losses = train(model, np.vstack([x, noise1_x]), np.vstack([y, noise1_y]),
               n_iter, lr)
plt.xlabel('iter')
plt.ylabel('loss')
plt.plot(range(n_iter), losses, "-")
plt.savefig("train2.jpg")
plt.cla()
print("test in data with two noise point")
test(model, noise2_x, noise2_y)